#!/usr/bin/env python

"""
This script is used to plot healpix map.
"""

import os, sys
import healpy as hp
import numpy as np
import matplotlib.pylab as plt
from argparse import ArgumentParser

def getTitle(heal_map_filename):
    if heal_map_filename.find('/') != -1:
        s = heal_map_filename.split('/')
        s = s[len(s)-1]
    else:
        s = heal_map_filename

    return s

def plot(heal_map, out_dir, min, max, norm, unit, coord, cmap, cbar, title):
    m = hp.read_map(args.heal_map)

    # idx = m < args.min
    # m[idx] = -1
    idx = m > args.max
    m[idx] = args.max

    title=getTitle(args.heal_map)
    hp.mollview( m, min=args.min, \
                    max=args.max, \
                    norm=args.norm, \
                    unit=args.unit, \
                    coord=args.coord, \
                    cmap=args.cmap, \
                    cbar=args.cbar,
                    title=getTitle(args.heal_map) )
    hp.graticule()
    title=title.replace('fits','png')

    if args.save_fig is True:
        plt.savefig(title, format='png')

    plt.show()

def hplot():
    parser = ArgumentParser()
    parser.add_argument('heal_map', type=str)
    parser.add_argument('--min', default=1, type=int)
    parser.add_argument('--max', default=1024, type=int)
    parser.add_argument('--norm', '-n', default='none', type=str, 
                        help='norm can be one of: none(linear), hist, log')
    parser.add_argument('--unit', '-u', default='visits', type=str)
    parser.add_argument('--coord', '-c', nargs='+', default='E', type=str,
                        help='E: Ecliptical; G: Galactic; C: Equatorial')
    # parser.add_argument('--cmap', default='nipy_spectral', type=str)
    parser.add_argument('--cmap', '-m', default='CMRmap', type=str)
    parser.add_argument('--cbar', '-b', action='store_true')
    parser.add_argument('--save_fig', '-s', action='store_true')

    args = parser.parse_args()

    if os.path.isfile(args.heal_map) is False:
        print('> {} does not exist'.format(args.heal_map))
        sys.exit(0)

    m = hp.read_map(args.heal_map)

    # idx = m < args.min
    # m[idx] = -1
    # idx = m > args.max
    # m[idx] = args.max

    m_max = m.max()

    coord = ['E']  # default
    # print('args.coord: {}'.format(args.coord))
    if args.coord[0] == 'C':
        print('> transforming from Ecliptcical to Equatorial')
        coord = ['E', 'C']
    if args.coord[0] == 'G':
        print('> transforming from Ecliptcical to Galactic')
        coord = ['E', 'G']

    title=getTitle(args.heal_map)
    hp.mollview( m, min=args.min, \
                    max=m_max, \
                    norm=args.norm, \
                    unit=args.unit, \
                    # coord=args.coord, \
                    coord=coord, \
                    cmap=args.cmap, \
                    cbar=args.cbar,
                    title=getTitle(args.heal_map) )
    hp.graticule(dpar=30,dmer=30,coord=args.coord)

    # hp.mollview( m, min=args.min, \
    #                 norm=args.norm, \
    #                 unit=args.unit, \
    #                 coord=args.coord, \
    #                 cmap=args.cmap, \
    #                 cbar=args.cbar,
    #                 sub=(2,2,2) )
    # hp.graticule()

    title=title.replace('fits','png')

    if args.save_fig is True:
        plt.savefig(title, format='png')

    plt.show()


if __name__ == '__main__':
    hplot()
